-module(mysql).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]).
-export([connect/5, connect/4, close/1, autocommit/2, set_charset/2, escape/1]).
-export([execute/2, execute/3, start_transaction/1, commit/1, rollback/1]).
-include("mysql_define.hrl").
-behaviour(gen_server).

-define(BACKSLASH, 92).
-define(QUOTE, $').
-define(DOUBLEQUOTE, $").

%%---------------------------------------------------------------------------
%% interface function
%%---------------------------------------------------------------------------

%%---------------------------------------------------------------------------
%% Function: connect(Host, Port, User, Passwd, Db) -> {ok, Pid} |
%%                                                    {error, Code, Message}
%% Description: Connect to mysql server without default db
%%---------------------------------------------------------------------------
connect(Host, Port, User, Passwd) ->
	connect(Host, Port, User, Passwd, undefined).

%%---------------------------------------------------------------------------
%% Function: connect(Host, Port, User, Passwd, Db) -> {ok, Pid} |
%%                                                    {error, Code, Message}
%% Description: Connect to mysql server with default db
%%---------------------------------------------------------------------------
connect(Host, Port, User, Passwd, Db) ->
	{ok, Pid} = gen_server:start_link(?MODULE, [], []),
	Reply = gen_server:call(Pid, {connect, Host, Port, User, Passwd, Db}),
	case Reply of
		ok ->
			{ok, Pid};
		Err ->
			Err
	end.

%%---------------------------------------------------------------------------
%% Function: execute(Pid, Statement) -> {ok, AffectedRow, InsertId} |
%%                                      {result, MetaData, Rows} |
%%                                      {error, Code, Message}
%% Description: send query to mysql server and fetch result
%%---------------------------------------------------------------------------
execute(Pid, Statement) ->
	gen_server:call(Pid, {execute, list_to_binary(Statement)}).


%%---------------------------------------------------------------------------
%% Function: execute(Pid, Statement, Parameters) -> {ok, AffectedRow, InsertId} |
%%                                      {result, MetaData, Rows} |
%%                                      {error, Code, Message}
%% Description: send query to mysql server and fetch result, with parameters bind
%%---------------------------------------------------------------------------
execute(Pid, Statement, Parameters) ->
	execute(Pid, Statement, Parameters, []).

execute(Pid, Statement, [], Result) ->
	L1 = [Statement|Result],
	L2 = lists:reverse(L1),
	S = iolist_to_binary(L2),
	io:format("~p~n", [S]),
	gen_server:call(Pid, {execute, S});
execute(Pid, Statement, [P|Parameters], Result) ->
	case string:chr(Statement, 63) of
		0 ->
			throw(badarg);
		Index ->
			Before = string:substr(Statement, 1, Index - 1),
			After = string:substr(Statement, Index + 1),
			Escaped = escape(P),
			L1 = [Before|Result],
			L2 = [Escaped|L1],
			execute(Pid, After, Parameters, L2)
	end.

%%---------------------------------------------------------------------------
%% Function: escape(String) -> EscapedString
%% Types:
%% 		String = string()
%% 		EscapedString = string()
%% Description: send query to mysql server and fetch result, with parameters bind
%%---------------------------------------------------------------------------
escape(Value) ->
	case Value of
		Value when is_list(Value) ->
			escape_string(Value, []);
		Value when is_integer(Value) ->
			list_to_binary(integer_to_list(Value));
		Value when is_binary(Value) ->
			escape_string(binary_to_list(Value), []);
		_ ->
			throw(badarg)
	end.
	
escape_string([], To) -> lists:reverse(To);
escape_string([0|From], To) -> L = [?BACKSLASH|To], escape_string(From, [0|L]);
escape_string([?QUOTE|From], To) -> L = [?BACKSLASH|To], escape_string(From, [?QUOTE|L]);
escape_string([?DOUBLEQUOTE|From], To) -> L = [?BACKSLASH|To], escape_string(From, [?DOUBLEQUOTE|L]);
escape_string([C|From], To) -> escape_string(From, [C|To]).

autocommit(Pid, Bool) ->
	case Bool of
		true ->
			gen_server:call(Pid, {execute, <<"set autocommit=on">>});
		false ->
			gen_server:call(Pid, {execute, <<"set autocommit=off">>});
		_ ->
			throw(badarg)
	end.

start_transaction(Pid) ->
	gen_server:call(Pid, {execute, <<"begin">>}).

commit(Pid) ->
	gen_server:call(Pid, {execute, <<"commit">>}).

rollback(Pid) ->
	gen_server:call(Pid, {execute, <<"rollback">>}).

set_charset(Pid, Charset) ->
	Statement = iolist_to_binary(["set names '", Charset, "'"]),
	gen_server:call(Pid, {execute, Statement}).

%%---------------------------------------------------------------------------
%% Function: close(Pid) -> stopped
%% Description: close connection
%%---------------------------------------------------------------------------
close(Pid) ->
	gen_server:call(Pid, {close}).

%%---------------------------------------------------------------------------
%% gen_server behaviour
%%---------------------------------------------------------------------------
init(_) ->
	{ok, state}.

handle_call({connect, Host, Port, User, Passwd, Db}, _From, State) ->
	case gen_tcp:connect(Host, Port, [binary, {active, false}]) of
		{ok, Sock} ->
			mysql_protocol:init(Sock),
			Reply = case Db of
				undefined ->
					login(list_to_binary(User), list_to_binary(Passwd), undefined);
				_ ->
					login(list_to_binary(User), list_to_binary(Passwd), list_to_binary(Db))
			end,
			{reply, Reply, State};
		_ ->
			{error, {2003, <<"#00000">>, "Can't connect to MySQL server on '127.0.0.1' (111)"}}
	end;
handle_call({execute, Statement}, _From, State) ->
	B = iolist_to_binary(Statement),
	Packet = <<?COM_QUERY:?UINT8, B/binary>>,
	mysql_protocol:send(0, Packet),

	Reply = case receive_result() of
		{ok, _FieldCount, AffectedRows, InsertId} -> %update ok
			{ok, {AffectedRows, InsertId}};
		{field_count, FieldCount} -> %query ok
			{MetaData, Rows} = fetch_result_set(FieldCount),
			{result, {MetaData, Rows}};
        {error, ErrorNumber, ErrorMessage} -> %execute error
			{error, {ErrorNumber, ErrorMessage}}
	end,
	{reply, Reply, State};
handle_call({close}, _From, State) ->
	{stop, normal, stopped, State}.

handle_cast(_Msg, State) ->  
    {noreply, State}. 

handle_info(_Info, State) ->  
    {noreply, State}.

terminate(_Reason, _State) -> ok.  

code_change(_OldVsn, State, _Extra) -> {ok, State}.

%%---------------------------------------------------------------------------
%% internal function
%%---------------------------------------------------------------------------
fetch_result_set(FieldCount) ->
	MetaData = fetch_fields(FieldCount, 0, []),
	{eof} = receive_result(),
	Rows = fetch_rows(MetaData, []),
	{MetaData, Rows}.

fetch_rows(MetaData, Rows) ->
	case mysql_protocol:recv() of
		{ok, _, Packet} ->
			<<Tag:?UINT8, _/binary>> = Packet,
			case Tag of
				254 ->	% end of rows
					lists:reverse(Rows);
				_ -> % more rows
					Row = fetch_columns(Packet, MetaData, []),
					fetch_rows(MetaData, [Row|Rows])
			end;
		Err ->
			Err
	end.

convert_column(Type, B) ->
	case Type of
		?MYSQL_TYPE_NULL -> null;
		?MYSQL_TYPE_TINY -> list_to_integer(binary_to_list(B));
		?MYSQL_TYPE_SHORT -> list_to_integer(binary_to_list(B));
		?MYSQL_TYPE_LONG -> list_to_integer(binary_to_list(B));
		?MYSQL_TYPE_LONGLONG -> list_to_integer(binary_to_list(B));
		?MYSQL_TYPE_VARCHAR -> binary_to_list(B);
		?MYSQL_TYPE_VAR_STRING -> binary_to_list(B);
		?MYSQL_TYPE_STRING -> binary_to_list(B);
		_ -> B
	end.

fetch_columns(_Buf, [], Columns) -> lists:reverse(Columns);
fetch_columns(Buf, [Field|MetaData], Columns) -> 
	{B, Buf1} = mysql_protocol:get_binary(Buf),
	{_, _, Type} = Field,

	Column = convert_column(Type, B),
	fetch_columns(Buf1, MetaData, [Column|Columns]).

fetch_fields(FieldCount, FieldCount, FieldList) -> lists:reverse(FieldList);
fetch_fields(FieldCount, Index, FieldList) ->
	{ok, _, Packet} = mysql_protocol:recv(),

	%decode field
	{_Catalog, Rest} = mysql_protocol:get_binary(Packet),
	{_Database, Rest2} = mysql_protocol:get_binary(Rest),
	{_Table, Rest3} = mysql_protocol:get_binary(Rest2),
	{_OrgTable, Rest4} = mysql_protocol:get_binary(Rest3),
	{Name, Rest5} = mysql_protocol:get_binary(Rest4),
	{_OrgName, Rest6} = mysql_protocol:get_binary(Rest5),
	<<_Metadata:8/little, _Charset:16/little, Length:32/little, Type:8/little, _Flags:16/little, _Decimals:8/little, _Rest7/binary>> = Rest6,

	Field = {binary_to_list(Name), Length, Type},
	fetch_fields(FieldCount, Index + 1, [Field|FieldList]).

receive_result() ->
	{ok, _Id, Packet} = mysql_protocol:recv(),
   	<<FieldCount:?UINT8, B0/binary>> = Packet,

	case FieldCount of
		255 -> % error
			<<ErrorNumber:?UINT16, _SqlState:6/binary, ErrorMessage/binary>> = B0,
			{error, ErrorNumber, binary_to_list(ErrorMessage)};
		254 -> % end of result set
			<<_WarningCount:?UINT16, _StatusFlags:?UINT16>> = B0,
			{eof};        
		0 -> % update success,
			{AffectedRows, B1} = mysql_protocol:get_length(B0),
			{InsertId, B2} =mysql_protocol:get_length(B1),
			<<_ServerStatus:?UINT16, _WarningCount:?UINT16, _Message/binary>> = B2,
			{ok, FieldCount, AffectedRows, InsertId};
		_ -> % query success
			{field_count, FieldCount}
	end.

login(User, Passwd, Db) ->
	% decode handshake packet
	{ok, 0, Handshake} = mysql_protocol:recv(),
	<<_ProtocolVersion:?UINT8, Rest1/binary>> = Handshake,
	{_ServerVersion, Rest2} = mysql_protocol:get_string(Rest1),

	<<_ThreadId:?UINT32,
	Scramble1:8/binary,
	_:1/binary,
	_ServerCapabilities1:?UINT16,
	_ServerLanguage:?UINT8,
	_ServerStatus:?UINT16,
	_ServerCapabilities2:?UINT16,
	_:11/binary,
	Scramble2 : 12/binary,
	_:1/binary>> = Rest2,

	Salt = iolist_to_binary([Scramble1,Scramble2]),
	_ServerCapabilities = _ServerCapabilities2 * 65536 + _ServerCapabilities1,

	% send authentication packet
	Flags = ?CLIENT_LONG_PASSWORD + ?CLIENT_LONG_FLAG + ?CLIENT_LOCAL_FILES + ?CLIENT_PROTOCOL_41 + ?CLIENT_SECURE_CONNECTION,
	Filler = <<0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0>>,
	MaxPacketSize = 16777216, % packet size up to 16M
	CharsetNumber = 10,
	ScramblePasswd = mysql_protocol:new_password(Passwd, Salt),
	PasswdLen = size(ScramblePasswd),

	case Db of
		undefined ->
			ClientFlags = Flags,
			Auth = <<ClientFlags:?UINT32,
				MaxPacketSize:?UINT32, 
				CharsetNumber:?UINT8, 
				Filler/binary, 
				User/binary, 0,
				PasswdLen, ScramblePasswd/binary>>;
		_ ->
			ClientFlags = Flags + ?CLIENT_CONNECT_WITH_DB,
			Auth = <<ClientFlags:?UINT32,
				MaxPacketSize:?UINT32,
				CharsetNumber:?UINT8,
				Filler/binary,
				User/binary, 0,
				PasswdLen, ScramblePasswd/binary,
				Db/binary, 0>>
	end,
	mysql_protocol:send(1, Auth),

	case receive_result() of
		{ok, _, _, _} ->
			ok;
		Err ->
			Err
	end.	

